package com.ruoyi.ai.service.impl;

import com.ruoyi.ai.dto.ChatRequestDTO;
import com.ruoyi.ai.service.ChatService;
import com.ruoyi.ai.service.VectorService;
import com.ruoyi.ai.tools.ChatTools;
import jakarta.annotation.PostConstruct;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.model.tool.ToolCallingChatOptions;
import org.springframework.ai.model.tool.ToolCallingManager;
import org.springframework.ai.model.tool.ToolExecutionResult;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.tool.definition.ToolDefinition;
import org.springframework.ai.tool.method.MethodToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.ReflectionUtils;
import reactor.core.publisher.Flux;

import java.lang.reflect.Method;
import java.util.List;
import java.util.Map;

import static com.ruoyi.ai.filter.HeaderFilter.DETAILS_USER_ID;


/**
 * @author 15332
 */
@Slf4j
@Service
public class ChatServiceImpl implements ChatService {

  /**
   * 提示词
   */
  public static final String PROMPT = """
    你是一个雅橙科技教学助手，由龙哥研制，正在回答用户的提问。
    所有的回答都只能通过中文回答。
    所有回答的格式都是纯文本格式。
    我们公司的口号是"用心做教育"，你只能回答编程问题以及教学相关的业务，。
    请使用toolcall查询的数据进行回复，如果没有查询到的业务数据，只需说明你不知道答案即可。
    历史对话：
    {history}


    问题
    {input}

    """;
  public static final double TEMPERATURE = 0.75;
  public static final double FREQUENCY_PENALTY = 1.5;
  public static final int MAX_TOKENS = 200;
  @Autowired
  private OpenAiChatModel chatModel;
  @Autowired
  private ChatTools chatTools;
  @Autowired
  private VectorService vectorService;

  // 创建工具回调集合
  private ToolCallback[] toolCallbacksSet;

  // 构建聊天选项
  private ChatOptions chatOptions;
  private ChatMemory chatMemory;


  @PostConstruct
  public void init() {
    toolCallbacksSet = new ToolCallback[]{
      createToolCallback(chatTools, "getCurrentDateTime", "获取用户当前的日期信息"),
      createToolCallback(chatTools, "queryTeachingPlan", "查询课表"),
      createToolCallback(chatTools, "route", "跳转页面"),
      createToolCallback(chatTools, "queryAssistantName", "查询助手信息"),
      createToolCallback(chatTools, "queryUserName", "查询用户信息和学习情况"),
      createToolCallback(chatTools, "queryCourse", "查询课程信息")};

    chatMemory = MessageWindowChatMemory
      .builder()
      .maxMessages(100)
      .chatMemoryRepository(new InMemoryChatMemoryRepository())
      .build();
  }

  @Override
  public Flux<ChatResponse> chatWithTools(ChatRequestDTO chatRequestDTO) {
    return Flux.deferContextual(ctx -> {

      String message = chatRequestDTO.message();
      boolean useHistoryChat = chatRequestDTO.useHistoryChat();
      boolean useTools = chatRequestDTO.useTools();
      boolean useLocalKnowledge = chatRequestDTO.useLocalKnowledge();


      ToolCallingManager toolCallingManager = ToolCallingManager.builder().build();
      String userId = ctx.get(DETAILS_USER_ID);
      //获取当前用户的聊天记录
      List<Message> historyMessages = chatMemory.get(userId);
      List<String> contentList = vectorService.searchQuestion(message);

      // 构建提示并发送请求
      Map<String, Object> promptParams =new java.util.HashMap<>(Map.of(
          "input", message,
          "history", historyMessages,
          "documents", contentList
        ));
      StringBuilder promptStr = new StringBuilder(PROMPT);
      //获取历史对话
      if (useHistoryChat) {
        historyMessages = chatMemory.get(userId);
        log.info("获取历史对话记录, 数量: {}", historyMessages.size());
        if (CollectionUtils.isNotEmpty(historyMessages)) {
          promptParams.put("history", historyMessages);
          promptStr.append("""
                                历史对话：
                                {history}
                            """);
        }
      }
      //获取当本地知识库文档
      if (useLocalKnowledge) {
        log.info("开始从向量数据库搜索相关文档");
        contentList = vectorService.searchQuestion(message);
        if (CollectionUtils.isNotEmpty(contentList)) {
          log.info("向量搜索完成, 获取到文档数量: {}", contentList.size());
          promptParams.put("documents", contentList);
          promptStr.append("""
                            文档:
                            {documents}
                            """);
        }
      }
      // 构建工具回调集合
      if (useTools) {
        chatOptions = ToolCallingChatOptions.builder()
          .toolCallbacks(toolCallbacksSet)
          .internalToolExecutionEnabled(false)
          .maxTokens(1000)
          .frequencyPenalty(1.5)
          .temperature(0.75)
          .build();
      }

      //获取当前用户信息
      // 创建提示模板
      PromptTemplate promptTemplate = new PromptTemplate(promptStr.toString());
      Prompt prompt = promptTemplate.create(promptParams, chatOptions);
      ChatResponse chatResponse = chatModel.call(prompt);

      // 处理工具调用
      while (chatResponse.hasToolCalls()) {
        ToolExecutionResult toolExecutionResult = toolCallingManager.executeToolCalls(prompt, chatResponse);
        prompt = new Prompt(toolExecutionResult.conversationHistory(), chatOptions);
        chatResponse = chatModel.call(prompt);
      }

      StringBuilder sb = new StringBuilder();
      return chatModel.stream(prompt).doOnNext(cr -> {
        //追加响应
        String text = cr.getResults().get(0).getOutput().getText();
        sb.append(text);
        log.debug("接收到流式响应片段: {}", text);
      }).doOnComplete(() -> {
        addHistoryMessage(sb, userId, message);
      });
    });

  }

  /**
   * 添加历史消息
   *
   * @param sb
   * @param userId
   * @param message
   */
  private void addHistoryMessage(StringBuilder sb, String userId, String message) {
    log.info("流式响应接收完成, 最终响应长度: {}", sb.length());
    //更新历史记录
    log.info("更新对话历史记录 - 添加用户消息");
    chatMemory.add(userId, new UserMessage(message));
    log.info("更新对话历史记录 - 添加助手响应");
    chatMemory.add(userId, new AssistantMessage(sb.toString()));
    log.info("对话历史记录更新完成");
  }

  /**
   * 创建工具回调
   *
   * @param toolObject  工具对象
   * @param methodName  方法名
   * @param description 描述
   * @return 工具回调
   */
  private ToolCallback createToolCallback(Object toolObject, String methodName, String description) {

    Method method = ReflectionUtils.findMethod(toolObject.getClass(), methodName);
    ToolDefinition toolDefinition = ToolDefinition
      .builder()
      .inputSchema("{}")
      .name(methodName)
      .description(description)
      .build();
    return MethodToolCallback.builder()
      .toolDefinition(toolDefinition)
      .toolMethod(method)
      .toolObject(toolObject)
      .build();
  }
}
